#include "defaults.h"
#include "utils.h"
#include "fe.h"
#include "math/numerics.h"
#include "wrapper/superluwrapper.h"
#ifdef _OPENMP
#include <omp.h>
#endif

using namespace std;

///////////////////////////////////////////////////////////////////////////////////////////////////////////////
void GlobalAssembly(const Element* e, const Matrix& A, map< pair<size_t, size_t>, nfloat >& entry) {
	for(size_t i = 0; i < A.size1(); i++) {
		for(size_t j = 0; j < A.size2(); j++) {
			entry[pair<size_t, size_t>(e->GetNode(i)->Id(), e->GetNode(j)->Id())] += A(i, j);
		}
	}
}

void GlobalAssemblyTri(const Element* e, const Matrix& A, map< pair<size_t, size_t>, nfloat >& entry) {
	for(size_t i = 0; i < A.size1(); i++) {
		const size_t ig = e->GetNode(i)->Id();
		for(size_t j = 0; j < A.size2(); j++) {
			const size_t jg = e->GetNode(j)->Id();
			if(jg >= ig)
				entry[pair<size_t, size_t>(ig, jg)] += A(i, j);
		}
	}
}

void GlobalAssembly(const Element* e, const Vector& blocal, Vector& bglobal) {
	for(size_t i = 0; i < e->SizeNodes(); i++) {
		const Node* n = e->GetNode(i);
		bglobal(n->Id()) += blocal(i);
	}
}
///////////////////////////////////////////////////////////////////////////////////////////////////////////////
Vector Force(Element* e, const std::string& attribute) {
	Vector b(e->SizeNodes());
	for(int i = 0; i < e->SizeNodes(); i++) {
		Node* n = e->GetNode(i);
		b(i) = e->Integrate<nfloat>(n, attribute);
	}
	return b;
}

SparseMatrix StiffnessMatrix(Mesh& mesh, const std::string& attribute, const bool triangular) {
	unsigned n = (unsigned)mesh.SizeNodes();
	map< pair<size_t, size_t>, nfloat > entry;
	size_t elements = mesh.SizeElements();
#ifdef _OPENMP
#pragma omp parallel for
#endif
    for(int i = 0; i < elements; i++) {
        Element* e = mesh.GetElement(i);
        cerr << (size_t)(100.0 * ((nfloat)i / (elements - 1))) << "%" << '\r';
        Matrix K = e->Stiffness(attribute);
#ifdef _OPENMP
#pragma omp critical
#endif
        {
			if(triangular)
				GlobalAssemblyTri(e, K, entry);
			else
				GlobalAssembly(e, K, entry);
        }
    }
	SparseMatrix A(n, n, entry.size());
	for(map< pair<size_t, size_t>, nfloat>::iterator it = entry.begin(); it != entry.end(); it++) {
		const pair< pair<size_t, size_t>, nfloat>& p = *it;
		const size_t& i = p.first.first;
		const size_t& j = p.first.second;
		const nfloat& a = p.second;
		A(i, j) = a;
	}

    return A;
}

SparseMatrix MassMatrix(Mesh& mesh, const std::string& attribute, const bool triangular, const bool lumped) {
	unsigned n = (unsigned)mesh.SizeNodes();
	map< pair<size_t, size_t>, nfloat > entry;
	size_t elements = mesh.SizeElements();
#ifdef _OPENMP
#pragma omp parallel for
#endif
    for(int i = 0; i < elements; i++) {
        Element* e = mesh.GetElement(i);
        cerr << (size_t)(100.0 * ((nfloat)i / (elements - 1))) << "%" << '\r';
        Matrix K = e->Mass(attribute, lumped);
#ifdef _OPENMP
#pragma omp critical
#endif
        {
			if(triangular)
				GlobalAssemblyTri(e, K, entry);
			else
				GlobalAssembly(e, K, entry);
        }
    }
	SparseMatrix A(n, n, entry.size());
	for(map< pair<size_t, size_t>, nfloat>::iterator it = entry.begin(); it != entry.end(); it++) {
		const pair< pair<size_t, size_t>, nfloat>& p = *it;
		const size_t& i = p.first.first;
		const size_t& j = p.first.second;
		const nfloat& a = p.second;
		A(i, j) = a;
	}

    return A;
}

Vector Force(Mesh& mesh, const std::string& attribute) {
	unsigned n = (unsigned)mesh.SizeNodes();
	map< pair<size_t, size_t>, nfloat > entry;
	size_t elements = mesh.SizeElements();
	Vector f = ZeroVector(n);
#ifdef _OPENMP
#pragma omp parallel for
#endif
    for(int i = 0; i < elements; i++) {
        Element* e = mesh.GetElement(i);
        cerr << (size_t)(100.0 * ((nfloat)i / (elements - 1))) << "%" << '\r';
		Vector fi = Force(e, attribute);
#ifdef _OPENMP
#pragma omp critical
#endif
        {
			GlobalAssembly(e, fi, f);
        }
    }

    return f;
}

nfloat MeanValue(Node* v1, Node* v2, Node* vLeft, Node* vRight, const std::string& attribute = "position") {
	Vec3n e = v2->Get<Vec3n>(attribute) - v1->Get<Vec3n>(attribute);
	Vec3n eL = vLeft->Get<Vec3n>(attribute) - v1->Get<Vec3n>(attribute);
	Vec3n eR = vRight->Get<Vec3n>(attribute) - v1->Get<Vec3n>(attribute);

	nfloat d = length(e);
			
	normalize(e);
	normalize(eL);
	normalize(eR);
	
	nfloat alpha = std::acos(dot(e, eL));
	nfloat beta = std::acos(dot(e, eR));
	
	nfloat w = (std::tan(alpha * 0.5) + std::tan(beta * 0.5)) / d;
	return w;
}

SparseMatrix MeanValueWeights(STriangleMesh& tm, const std::string& attribute) {
	unsigned n = tm.GetNumVertices();
	
	map< pair<size_t, size_t>, nfloat > entry;
	for(size_t i = 0; i < n; i++) {
		cout << i << "/" << n - 1 << '\r';
		Node* v1 = tm.Vertex(i);

		nfloat sum = 0.0;
		size_t m = v1->NumNeighbors();
		for(size_t j = 0; j < m; j++) {
			int li = (j - 1 + m) % m;
			int ri = (j + 1) % m;
			Node* vLeft = v1->Neighbor(li);
			Node* vRight = v1->Neighbor(ri);
			Node* v2 = v1->Neighbor(j);
			const nfloat w = MeanValue(v1, v2, vLeft, vRight, attribute);
			
			entry[pair<size_t, size_t>(v1->Id(), v2->Id())] = w;
			
			sum += -w;
		}
			
				
		entry[pair<size_t, size_t>(v1->Id(), v1->Id())] = sum;
	}
	
// 	cout <<  entry.size() << "(" << n + tm.GetNumEdges() << ") non-zero entries." << endl;
	SparseMatrix A(n, n, entry.size());
	for(map< pair<size_t, size_t>, nfloat>::iterator it = entry.begin(); it != entry.end(); it++) {
		const pair< pair<size_t, size_t>, nfloat>& p = *it;
		A(p.first.first, p.first.second) = p.second;
	}

	return A;
}

nfloat CoTangent(Node* v1, Node* v2, Node* vLeft, Node* vRight, const std::string& attribute = "position") {
#define COT(alpha)	(1.0 / std::tan(alpha))
	Vec3n a1 = v1->Get<Vec3n>(attribute) - vRight->Get<Vec3n>(attribute);
	Vec3n a2 = v2->Get<Vec3n>(attribute) - vRight->Get<Vec3n>(attribute);
	normalize(a1);
	normalize(a2);
	
	Vec3n b1 = v1->Get<Vec3n>(attribute) - vLeft->Get<Vec3n>(attribute);
	Vec3n b2 = v2->Get<Vec3n>(attribute) - vLeft->Get<Vec3n>(attribute);
	normalize(b1);
	normalize(b2);

	const nfloat alpha = std::acos(dot(a1, a2));
	const nfloat beta = std::acos(dot(b1, b2));
	
	return COT(alpha) + COT(beta);
}

SparseMatrix CotangentWeights(STriangleMesh& tm, const std::string& attribute) {
	unsigned n = tm.GetNumVertices();
	
	map< pair<size_t, size_t>, nfloat > entry;
	for(size_t i = 0; i < n; i++) {
		cout << i << "/" << n - 1 << '\r';
		Node* v1 = tm.Vertex(i);

		nfloat sum = 0.0;
		size_t m = v1->NumNeighbors();
		for(size_t j = 0; j < m; j++) {
			int li = (j - 1 + m) % m;
			int ri = (j + 1) % m;
			Node* vLeft = v1->Neighbor(li);
			Node* vRight = v1->Neighbor(ri);
			Node* v2 = v1->Neighbor(j);
			const nfloat w = CoTangent(v1, v2, vLeft, vRight, attribute);
			
			entry[pair<size_t, size_t>(v1->Id(), v2->Id())] = w;
			
			sum += -w;
		}
			
				
		entry[pair<size_t, size_t>(v1->Id(), v1->Id())] = sum;
	}
	
// 	cout <<  entry.size() << "(" << n + tm.GetNumEdges() << ") non-zero entries." << endl;
	SparseMatrix A(n, n, entry.size());
	for(map< pair<size_t, size_t>, nfloat>::iterator it = entry.begin(); it != entry.end(); it++) {
		const pair< pair<size_t, size_t>, nfloat>& p = *it;
		A(p.first.first, p.first.second) = p.second;
	}

	return A;
}
